#!/usr/bin/env python

import sys
import copy
#import profile
#import aho_corasick2
import time
#import suffix_trie

from de.unibi.comet.ac import AhoCorasickAutomaton
from de.unibi.comet.ac import ExtAhoCorasickAutomaton
from de.unibi.comet.ac import StringSetProvider
from de.unibi.comet.util import Log
from java.util import ArrayList

def usage():
  print "usage: %s <score_file> <quality_file>" % sys.argv[0]
  sys.exit(1)

class PWM(object):
    def __init__(self, name):
        self.weight=[]
        self.name=name
        self.threshold=None
        self.accuracy=1000.0
    def __str__(self):
        s ='PWM "%s":\n'%self.name
        s+='threshold: %s\n'%str(self.threshold)
        s+=self.matrix2str(self.weight)
        return s      
    def matrix2str(self, m):
        s=' pos |   A    |   C    |   G    |   T    | max_abs|  min   |  max   |\n'
        for row in m:
            s+='%4d |'%row[1]
            for f in row[0]:
                s+="%+7d |" % f
            s+="%+7d |" % row[2] # max_abs
            s+="%+7d |" % row[3] # min
            s+="%+7d |" % row[4] # max
            s+='\n'
        return s
    def addRow(self, row):
        int_row=[int(round(f*self.accuracy)) for f in row]
        self.weight.append((int_row, len(self.weight), max([abs(f) for f in int_row]), min(int_row), max(int_row), max(int_row)-min(int_row)))
    def setThreshold(self, t):
        self.threshold=int(round(t*self.accuracy))
    def calcScoreRange(self, s):
        if len(s)!=len(self.weight): return None
        min_score=0.0
        max_score=0.0
        for k in xrange(len(s)):
            if s[k]=='?':
                min_score+=self.weight[k][3]
                max_score+=self.weight[k][4]
                continue
            char_idx=-1
            if s[k]=='a': char_idx=0
            if s[k]=='c': char_idx=1
            if s[k]=='g': char_idx=2
            if s[k]=='t': char_idx=3
            if char_idx!=-1:
                min_score+=self.weight[k][0][char_idx]
                max_score+=self.weight[k][0][char_idx]
        return (min_score/self.accuracy, max_score/self.accuracy)
    def generateStrings(self):
        def cmp_func(x,y):
            if x[2]<y[2]: return -1
            elif x[2]>y[2]: return 1
            else: return 0
            
        # workaround for pyhton 2.2 not having w_sorted.sort(cmp=cmp_func, reverse=True)
        l=[]
        for w in self.weight:
            l.append((w[2],w))
        l.sort()
        w_sorted=[]
        for i in xrange(len(l)-1,-1,-1):
            w_sorted.append(l[i][1])
#        print 'sorted:\n', self.matrix2str(w_sorted)
        chars='acgt'
        result_list=[]
        def generate_rec(s, idx, score, min_score_left, max_score_left):
#            print 'generate_rec: s=%s, idx=%d, score=%.3f, min_score_left=%.3f, max_score_left=%.3f'%\
#            (s, idx, score, min_score_left, max_score_left)
            if len(result_list)>=50000: return
            if score+min_score_left>=self.threshold:
#                print '  storing string'
                result_list.append(s)
                return
            if score+max_score_left<self.threshold:
#                print '  abandoning path'
                return
#            if (idx==len(w_sorted)) and (score>=self.threshold):
#                result_list.append(s)
            if idx>=len(w_sorted): return
            pos=w_sorted[idx][1]
            new_min_score_left=min_score_left-w_sorted[idx][3]
            new_max_score_left=max_score_left-w_sorted[idx][4]
            for k in xrange(4):
                s_new=s[0:pos]+chars[k]+s[pos+1:]
                generate_rec(s_new, idx+1, score+w_sorted[idx][0][k], new_min_score_left,new_max_score_left)
        total_min_score=sum([row[3] for row in w_sorted])
        total_max_score=sum([row[4] for row in w_sorted])
        s=''.join(len(w_sorted)*['?'])
        generate_rec(s,0,0,total_min_score,total_max_score)
        return result_list
    # generate strings without wildcards
    def generateAllStrings(self):
        w_sorted=copy.copy(self.weight)
        chars='acgt'
        result_list=[]
        def generate_rec(s, idx, score, min_score_left, max_score_left):
            if len(result_list)>100000: return
            if score+max_score_left<self.threshold:
                return
            if (idx==len(w_sorted)) and (score>=self.threshold):
                result_list.append(s)
            if idx>=len(w_sorted): return
            pos=w_sorted[idx][1]
            new_min_score_left=min_score_left-w_sorted[idx][3]
            new_max_score_left=max_score_left-w_sorted[idx][4]
            for k in xrange(4):
                s_new=s[0:pos]+chars[k]+s[pos+1:]
                generate_rec(s_new, idx+1, score+w_sorted[idx][0][k], new_min_score_left,new_max_score_left)
        total_min_score=sum([row[3] for row in w_sorted])
        total_max_score=sum([row[4] for row in w_sorted])
        s=''.join(len(w_sorted)*['?'])
        generate_rec(s,0,0,total_min_score,total_max_score)
        return result_list
    def countStrings(self):
        def cmp_func(x,y):
            if x[2]<y[2]: return -1
            elif x[2]>y[2]: return 1
            else: return 0
        w_sorted=copy.copy(self.weight)
        w_sorted.sort(cmp=cmp_func, reverse=True)
        # print 'sorted:\n', self.matrix2str(w_sorted)
        chars='acgt'
        count=[0,0]
        def count_rec(idx, score, min_score_left, max_score_left):
            if count[0]>=1000000: return
            if score+min_score_left>=self.threshold:
                count[0]+=1
                count[1]+=4**(len(w_sorted)-idx)
                return
            if score+max_score_left<self.threshold:
                return
            if idx>=len(w_sorted): return
            pos=w_sorted[idx][1]
            new_min_score_left=min_score_left-w_sorted[idx][3]
            new_max_score_left=max_score_left-w_sorted[idx][4]
            for k in xrange(4):
                count_rec(idx+1, score+w_sorted[idx][0][k], new_min_score_left,new_max_score_left)
        total_min_score=sum([row[3] for row in w_sorted])
        total_max_score=sum([row[4] for row in w_sorted])
        count_rec(0,0,total_min_score,total_max_score)
        return (count[0], count[1])

def expand(l):
    l_exp = []
    def rec(s,p):
#        print 'expand: rec: s=%s, p=%d'%(s,p)
        i = s.find('?', p)
        if i>=0:
            rec(s[0:i]+'a'+s[i+1:], i+1)
            rec(s[0:i]+'c'+s[i+1:], i+1)
            rec(s[0:i]+'g'+s[i+1:], i+1)
            rec(s[0:i]+'t'+s[i+1:], i+1)
        else:
            l_exp.append(s)
    for s in l:
        rec(s,0)
    return l_exp
        

def read_pwm_score_file(f):
    pwms={}
    pwm=None
    for line in f.readlines():
        #print "line: %s"%line,
        if line.startswith('>'):
            k=line.find('|')
            if k==-1: pwm=PWM(line[1:])
            else: pwm=PWM(line[1:k])
            continue
        if line.startswith('<') and (pwm!=None):
            if pwms.has_key(pwm.name):
                print 'warning: overwriting pwm "%s"'%pwm.name
            pwms[pwm.name]=pwm
            pwm=None
            continue
        if pwm!=None:
            fields=line.split()
            row = [float(field) for field in fields]
            if len(row)==4:
                pwm.addRow(row)
            else:
                print 'warning: in pwm "%s": ignoring row with %d entries'%(pwm.name,len(row))
    return pwms

def read_pwm_quality_file(f,pwms):
    name=''
    for line in f.readlines():
        if line.startswith('#'):
            if line.count('|')==2:
                name=line[2:].split('|')[0]
        else:
            fields = line.split()
            if (len(fields)>=4) and (name!=''):
                if pwms.has_key(name):
                    pwms[name].setThreshold(float(fields[3]))
                else:
                    print 'warning: pwm "%s" not found (reading quality file)'

def is_minimal(l):
    n=len(l[0])
    for k in xrange(n):
#        print 'considering k=%d'%k
        l2=[]
        for s in l: l2.append(s[0:k]+s[k+1:]+s[k])
        l2.sort()
        def equal(i,j):
            if (i<0) or (i>=len(l2)) or (j<0) or (j>=len(l2)): return False
            for m in xrange(n-2,-1,-1):
                if l2[i][m]!=l2[j][m]: return False
            return True
        m=0
        for i in xrange(len(l2)):
            if equal(i-1,i): m+=1
            else: m=0
            if m==3:
                print 'k=%d'%k
                for h in range(i-3,i+1):
                    print l2[h]

def main():
#    if len(sys.argv)!=3: usage()
#    f_score=open(sys.argv[1])
#    f_quality=open(sys.argv[2])

    f_score=open('/homes/tmarscha/svn/diplom/data/pwm.all.score-50')
    f_quality=open('/homes/tmarscha/svn/diplom/data/pwm.all.quality-50')
    pwms = read_pwm_score_file(f_score)
    read_pwm_quality_file(f_quality,pwms)
    
    #for pwm in pwms.itervalues():
    #    l=pwm.generateStrings()
    #    print '"%s", (%d strings):'%(pwm.name,len(l))
    #    is_minimal(l)
    #    sys.stdout.flush()
    #    print pwm
    
#    pwm=pwms['P$ATHB9_01']  # 109605, 581136
#    pwm=pwms['V$FOXO4_02'] # 3294, 17370
#    pwm=pwms['P$CPRF2_Q2']  # 1692
    pwm=pwms['V$MEIS1_01']  # 28, 34432
    
#    for s in expand(pwm.generateStrings()):
#        print s
#    sys.exit(1)
    
    #pwm=pwms['V$SRY_01']
    #print pwm
    
    #count=0
    #total_count=0
    #for pwm in pwms.values():
    #    t=pwm.countStrings()
    #    count+=t[0]
    #    total_count+=t[1]
    #    print '"%s": '%pwm.name, t
    #    sys.stdout.flush()
    #    
    #print 'ratio=%f'%float(count)/float(total_count)
    
#    for s in pwm.generateAllStrings():
#        print s
#    sys.exit(1)

    Log.getInstance().setTimingActive(True)

    for pwm in pwms.itervalues():
        print '------ pwm=%s, length=%d'%(pwm.name, len(pwm.weight))
        results=[]
        results.append(pwm.name)
        results.append(len(pwm.weight))
        l_wc = pwm.generateStrings()
        l = expand(l_wc)
        print '--- all strings (without wildcards): %d'%len(l)
        results.append(len(l))
        java_list = ArrayList(len(l))
        for s in l: java_list.add(s)
        ssp = StringSetProvider(java_list)
        aca = AhoCorasickAutomaton()
        aca.build(ssp)
        results.append(Log.getInstance().getLastPeriod())
        results.append(aca.countNodes())
        print 'nodes=%d'%aca.countNodes()
    
        l = l_wc
        print '--- wildcard strings: %d'%len(l)
        results.append(len(l))
        java_list = ArrayList(len(l))
        for s in l: java_list.add(s)
        ssp = StringSetProvider(java_list)
        ext_aca = ExtAhoCorasickAutomaton()
        ext_aca.build(ssp)
        results.append(Log.getInstance().getLastPeriod())
        results.append(ext_aca.countNodes())
        results.append(ext_aca.countWildcardNodes())    
        print 'nodes=%d, wildcard nodes=%d'%(ext_aca.countNodes(), ext_aca.countWildcardNodes())
        print '> ', results
    
    sys.exit(1)

if __name__ == "__main__":
    #profile.run('main()', 'ac_test2.prof')
    main()
